#include "nalu.h"


#define nalu_debug_print
#ifdef nalu_debug_print
#define nalu_debug(fmt,...) s264_print(fmt, ##__VA_ARGS__)
#else
#define nalu_debug(fmt,...)
#endif



/// @brief 去除防竞争码 00 00 30 (00 | 01 | 02 | 03) => 00 00 (00 | 01 | 02 | 03)
/// @param nostartcode_buf startcode后的数据
/// @param nostartcode_buf_len startcode后的数据长度
/// @param rbsp_buf s264buf_t 变量地址(无须初始化,但需要释放)
/// @return 
i8 nalu_rbsp(u8 *nostartcode_buf, u64 nostartcode_buf_len, s264buf_t *rbsp_buf)
{
	if (nostartcode_buf == NULL || nostartcode_buf_len <= 0)
		return -1;
	if (s264buf_init(rbsp_buf, nostartcode_buf_len))
		return -2;
	for (size_t i = 0; i < nostartcode_buf_len; i++)
	{
		// 防竞争码 3
		if (nostartcode_buf[i] == 3)
		{
			if (i > 2)
			{
				// 前2位均为0
				if (nostartcode_buf[i - 1] == 0 && nostartcode_buf[i - 2] == 0)
				{
					if (i < nostartcode_buf_len - 1)
					{
						// 后一位 为 0 | 1 | 2 | 3
						if (
							nostartcode_buf[i + 1] == 0 ||
							nostartcode_buf[i + 1] == 1 ||
							nostartcode_buf[i + 1] == 2 ||
							nostartcode_buf[i + 1] == 3
							)
						{
							continue;
						}

					}
				}
			}
		}
		s264buf_append(rbsp_buf, &(nostartcode_buf[i]), 1);
	}
	
	return 0;
}




void *nalu_parse(nalu_t *nalu, u8 *type)
{
	if (nalu == NULL || type == NULL)
		return NULL;
	void *res = NULL;
	s264buf_t rbsp = { 0 };
	if (nalu_rbsp(nalu->buf + nalu->startcode_len, nalu->buf_len - nalu->startcode_len, &rbsp) == 0)
	{
		u8 nalu_forbidden_bit = nalu_forbidden_bit(rbsp.data[0]);
		u8 nalu_nal_ref_idc = nalu_nal_ref_idc(rbsp.data[0]);
		u8 nalu_nal_unit_type = nalu_nal_unit_type(rbsp.data[0]);
		bs_t bs;
		bs_init(&bs, rbsp.data + 1, rbsp.len - 1);
		if (nalu_nal_unit_type == nalu_type_sps)
		{
			res = (void *)nalu_parse_sps(&bs);
		}
		if (nalu_nal_unit_type == nalu_type_pps)
		{
			// res = nalu_parse_pps(rbsp.data+1,rbsp.len-1);
		}
		s264buf_deinit(&rbsp);
		if (type != NULL)
			*type = nalu_nal_unit_type;
	}
	return res;
}

static inline void rbsp_trailing_bits(bs_t *bs)
{
	if (bs == NULL)
		return;
	u8 rbsp_stop_one_bit = bs_read1bit(bs); // equal 1
	nalu_debug("rbsp_stop_one_bit:%d", rbsp_stop_one_bit);
	while (bs_eof(bs) == 0)
	{
		u8 rbsp_alignment_zero_bit = bs_read1bit(bs); // equal 0
		nalu_debug("rbsp_alignment_zero_bit:%d", rbsp_alignment_zero_bit);
	}
}

static inline void nalu_sps_scaling_list(u8 *list, u64 size, u8 *use_default,bs_t *bs)
{
	u8 last_scale = 8;
	u8 next_scale = 8;
	for (size_t i = 0; i < size; i++)
	{
		if (next_scale != 0)
		{
			i8 delta_scale = bs_readse(bs);
			next_scale = (last_scale + delta_scale + 256) % 256;
			use_default[i] = (i == 0 && next_scale == 0);
		}
		list[i] = (next_scale == 0) ? last_scale : next_scale;
		last_scale = list[i];
	}
}
static void nalu_hrd(nalu_hrd_t *hrd, bs_t *bs)
{
	if (hrd == NULL || bs == NULL)
		return;
	hrd->cpb_cnt_minus1 = bs_readue(bs);
	nalu_debug("cpb_cnt_minus1:%ld", hrd->cpb_cnt_minus1);
	hrd->bit_rate_scale = bs_readnbit(bs, 4);
	nalu_debug("bit_rate_scale:%d", hrd->bit_rate_scale);
	hrd->cpb_size_scale = bs_readnbit(bs, 4);
	nalu_debug("cpb_size_scale:%d", hrd->cpb_size_scale);
	if (hrd->cpb_cnt_minus1 > 0)
	{
		hrd->bit_rate_value_minus1 = s264_malloc(sizeof(u64) * hrd->cpb_cnt_minus1);
		hrd->cpb_size_value_minus1 = s264_malloc(sizeof(u64) * hrd->cpb_cnt_minus1);
		hrd->cbr_flag = s264_malloc(sizeof(u64) * hrd->cpb_cnt_minus1);
	}
	for (u64 i = 0; i < hrd->cpb_cnt_minus1; i++)
	{
		u64 bit_rate_value_minus1 = bs_readue(bs);
		nalu_debug("bit_rate_value_minus1[%ld]:%ld",i,bit_rate_value_minus1);
		if (hrd->bit_rate_value_minus1 != NULL)
			hrd->bit_rate_value_minus1[i] = bit_rate_value_minus1;

		u64 cpb_size_value_minus1 = bs_readue(bs);
		nalu_debug("cpb_size_value_minus1[%ld]:%ld",i,cpb_size_value_minus1);
		if (hrd->cpb_size_value_minus1 != NULL)
			hrd->cpb_size_value_minus1[i] = cpb_size_value_minus1;
		u8 cbr_flag = bs_read1bit(bs);
		nalu_debug("cbr_flag[%ld]:%d",i,cbr_flag);
		if (hrd->cbr_flag != NULL)
			hrd->cbr_flag[i] = cbr_flag;
	}
	hrd->initial_cpb_removal_delay_length_minus1 = bs_readnbit(bs, 5);
	nalu_debug("initial_cpb_removal_delay_length_minus1:%d", hrd->initial_cpb_removal_delay_length_minus1);
	hrd->cpb_removal_delay_length_minus1 = bs_readnbit(bs, 5);
	nalu_debug("cpb_removal_delay_length_minus1:%d", hrd->cpb_removal_delay_length_minus1);
	hrd->dpb_output_delay_length_minus1 = bs_readnbit(bs, 5);
	nalu_debug("dpb_output_delay_length_minus1:%d", hrd->dpb_output_delay_length_minus1);
	hrd->time_offset_length = bs_readnbit(bs, 5);
	nalu_debug("time_offset_length:%d", hrd->time_offset_length);
}

static inline void nalu_vui(nalu_vui_t *vui,bs_t *bs)
{
	if (bs == NULL || vui == NULL)
		return ;
	
	vui->aspect_ratio_info_present_flag = bs_read1bit(bs);
	nalu_debug("aspect_ratio_info_present_flag:%d", vui->aspect_ratio_info_present_flag);

	if (vui->aspect_ratio_info_present_flag)
	{
		vui->aspect_ratio_idc = bs_readnbit(bs, 8);
		nalu_debug("aspect_ratio_idc:%d", vui->aspect_ratio_idc);
		if (vui->aspect_ratio_idc == 255)
		{
			vui->sar_width = bs_readnbit(bs, 16);
			nalu_debug("sar_width:%d", vui->sar_width);
			vui->sar_height = bs_readnbit(bs, 16);
			nalu_debug("sar_height:%d", vui->sar_height);
		}
	}
	vui->overscan_info_present_flag = bs_read1bit(bs);
	nalu_debug("overscan_info_present_flag:%d", vui->overscan_info_present_flag);
	if (vui->overscan_info_present_flag)
	{
		vui->overscan_appropriate_flag = bs_read1bit(bs);
		nalu_debug("overscan_appropriate_flag:%d", vui->overscan_appropriate_flag);
	}
	vui->video_signal_type_present_flag = bs_read1bit(bs);
	nalu_debug("video_signal_type_present_flag:%d", vui->video_signal_type_present_flag);
	if (vui->video_signal_type_present_flag)
	{
		vui->video_format = bs_readnbit(bs, 3);
		nalu_debug("video_format:%d", vui->video_format);
		vui->video_full_range_flag = bs_read1bit(bs);
		nalu_debug("video_full_range_flag:%d", vui->video_full_range_flag);
		vui->colour_description_present_flag = bs_read1bit(bs);
		nalu_debug("colour_description_present_flag:%d", vui->colour_description_present_flag);
		if (vui->colour_description_present_flag)
		{
			vui->colour_primaries = bs_readnbit(bs, 8);
			nalu_debug("colour_primaries:%d", vui->colour_primaries);
			vui->transfer_characteristics = bs_readnbit(bs, 8);
			nalu_debug("transfer_characteristics:%d", vui->transfer_characteristics);
			vui->matrix_coefficients = bs_readnbit(bs, 8);
			nalu_debug("matrix_coefficients:%d", vui->matrix_coefficients);
		}
	}
	vui->chroma_loc_info_present_flag = bs_read1bit(bs);
	nalu_debug("chroma_loc_info_present_flag:%d", vui->chroma_loc_info_present_flag);
	if (vui->chroma_loc_info_present_flag)
	{
		vui->chroma_sample_loc_type_top_field = bs_readue(bs);
		nalu_debug("chroma_sample_loc_type_top_field:%ld", vui->chroma_sample_loc_type_top_field);
		vui->chroma_sample_loc_type_bottom_field = bs_readue(bs);
		nalu_debug("chroma_sample_loc_type_bottom_field:%ld", vui->chroma_sample_loc_type_bottom_field);
	}
	vui->timing_info_present_flag = bs_read1bit(bs);
	nalu_debug("timing_info_present_flag:%d", vui->timing_info_present_flag);
	if (vui->timing_info_present_flag)
	{
		vui->num_units_in_tick = bs_readnbit(bs, 32);
		nalu_debug("num_units_in_tick:%d", vui->num_units_in_tick);
		vui->time_scale = bs_readnbit(bs, 32);
		nalu_debug("time_scale:%d", vui->time_scale);
		vui->fixed_frame_rate_flag = bs_read1bit(bs);
		nalu_debug("fixed_frame_rate_flag:%d", vui->fixed_frame_rate_flag);
	}
	vui->nal_hrd_parameters_present_flag = bs_read1bit(bs);
	nalu_debug("nal_hrd_parameters_present_flag:%d", vui->nal_hrd_parameters_present_flag);
	if (vui->nal_hrd_parameters_present_flag)
	{
		nalu_hrd(&(vui->hrd1),bs);
	}
	vui->vcl_hrd_parameters_present_flag = bs_read1bit(bs);
	nalu_debug("vcl_hrd_parameters_present_flag:%d", vui->vcl_hrd_parameters_present_flag);
	if (vui->vcl_hrd_parameters_present_flag)
	{
		nalu_hrd(&(vui->hrd2),bs);
	}
	if (vui->nal_hrd_parameters_present_flag || vui->vcl_hrd_parameters_present_flag)
	{
		vui->low_delay_hrd_flag = bs_read1bit(bs);
		nalu_debug("low_delay_hrd_flag:%d", vui->low_delay_hrd_flag);
	}
	vui->pic_struct_present_flag = bs_read1bit(bs);
	nalu_debug("pic_struct_present_flag:%d", vui->pic_struct_present_flag);
	vui->bitstream_restriction_flag = bs_read1bit(bs);
	nalu_debug("bitstream_restriction_flag:%d", vui->bitstream_restriction_flag);
};



nalu_sps_t *nalu_parse_sps(bs_t *bs)
{
	if (bs == NULL)
		return NULL;
	nalu_sps_t *res = (nalu_sps_t *) s264_malloc(sizeof(nalu_sps_t));
	if (res == NULL)
		return NULL;
	s264_memset(res, 0, sizeof(nalu_sps_t));

	res->profile_idc = bs_readnbit(bs, 8);
	nalu_debug("profile_idc:%d", res->profile_idc);
	res->constraint_set0_flag = bs_read1bit(bs);
	nalu_debug("constraint_set0_flag:%d", res->constraint_set0_flag);
	res->constraint_set1_flag = bs_read1bit(bs);
	nalu_debug("constraint_set1_flag:%d", res->constraint_set1_flag);
	res->constraint_set2_flag = bs_read1bit(bs);
	nalu_debug("constraint_set2_flag:%d", res->constraint_set2_flag);
	res->constraint_set3_flag = bs_read1bit(bs);
	nalu_debug("constraint_set3_flag:%d", res->constraint_set3_flag);
	res->reserved_zero_4bits = bs_readnbit(bs, 4);
	nalu_debug("reserved_zero_4bits:%d", res->reserved_zero_4bits);
	res->level_idc = bs_readnbit(bs, 8);
	nalu_debug("level_idc:%d", res->level_idc);
	res->seq_parameter_set_id = bs_readue(bs);
	nalu_debug("seq_parameter_set_id:%ld", res->seq_parameter_set_id);
	if (res->profile_idc == 100 ||
		res->seq_parameter_set_id == 110 ||
		res->seq_parameter_set_id == 122 ||
		res->seq_parameter_set_id == 144
		)
	{
		res->chroma_format_idc = bs_readue(bs);
		nalu_debug("chroma_format_idc:%ld", res->chroma_format_idc);
		if (res->chroma_format_idc == 3)
		{
			res->residual_colour_transform_flag = bs_read1bit(bs);
			nalu_debug("residual_colour_transform_flag:%d", res->residual_colour_transform_flag);
		}
		res->bit_depth_luma_minus8 = bs_readue(bs);
		nalu_debug("bit_depth_luma_minus8:%ld", res->bit_depth_luma_minus8);
		res->bit_depth_chroma_minus8 = bs_readue(bs);
		nalu_debug("bit_depth_chroma_minus8:%ld", res->bit_depth_chroma_minus8);
		res->qpprime_y_zero_transform_bypass_flag = bs_read1bit(bs);
		nalu_debug("qpprime_y_zero_transform_bypass_flag:%d", res->qpprime_y_zero_transform_bypass_flag);
		res->seq_scaling_matrix_present_flag = bs_read1bit(bs);
		nalu_debug("seq_scaling_matrix_present_flag:%d", res->seq_scaling_matrix_present_flag);
		if (res->seq_scaling_matrix_present_flag)
		{
			u8 scaling_list4x4[4 * 4] = { 0 };
			u8 scaling_list8x8[8 * 8] = { 0 };
			u8 use_default4x4[4 * 4] = { 0 };
			u8 use_default8x8[8 * 8] = { 0 };
			for (u8 i = 0; i < 8; i++)
			{
				res->seq_scaling_list_present_flag[i] = bs_read1bit(bs);
				nalu_debug("seq_scaling_list_present_flag[%d]:%d", i,res->seq_scaling_list_present_flag[i]);
				if (res->seq_scaling_list_present_flag[i])
				{
					if (i < 6)
					{
						nalu_sps_scaling_list(scaling_list4x4, 16, &(use_default4x4[i]), bs);
					}
					else
					{
						nalu_sps_scaling_list(scaling_list8x8, 64, &(use_default8x8[i]), bs);
					}
				}
			}
		}

	}
	res->log2_max_frame_num_minus4 = bs_readue(bs);
	nalu_debug("log2_max_frame_num_minus4:%ld", res->log2_max_frame_num_minus4);
	res->pic_order_cnt_type = bs_readue(bs);
	nalu_debug("pic_order_cnt_type:%ld", res->pic_order_cnt_type);
	if (res->pic_order_cnt_type == 0)
	{
		res->log2_max_pic_order_cnt_lsb_minus4 = bs_readue(bs);
		nalu_debug("log2_max_pic_order_cnt_lsb_minus4:%ld", res->log2_max_pic_order_cnt_lsb_minus4);
	}
	else if (res->pic_order_cnt_type == 1)
	{
		res->delta_pic_order_always_zero_flag = bs_read1bit(bs);
		nalu_debug("delta_pic_order_always_zero_flag:%d", res->delta_pic_order_always_zero_flag);
		res->offset_for_non_ref_pic = bs_readse(bs);
		nalu_debug("offset_for_non_ref_pic:%ld", res->offset_for_non_ref_pic);
		res->offset_for_top_to_bottom_field = bs_readse(bs);
		nalu_debug("offset_for_top_to_bottom_field:%ld", res->offset_for_top_to_bottom_field);
		res->num_ref_frames_in_pic_order_cnt_cycle = bs_readue(bs);
		nalu_debug("num_ref_frames_in_pic_order_cnt_cycle:%ld", res->num_ref_frames_in_pic_order_cnt_cycle);
		if (res->num_ref_frames_in_pic_order_cnt_cycle > 0)
		{
			res->offset_for_ref_frame = (i64 *) s264_malloc(sizeof(i64) * res->num_ref_frames_in_pic_order_cnt_cycle);
			for (u64 i = 0; i < res->num_ref_frames_in_pic_order_cnt_cycle; i++)
			{
				i64 offset_for_ref_frame = bs_readse(bs);
				nalu_debug("offset_for_ref_frame[%ld]:%ld", i, offset_for_ref_frame);
				if (res->offset_for_ref_frame != NULL)
					res->offset_for_ref_frame[i] = offset_for_ref_frame;
			}
		}
	}
	res->num_ref_frames = bs_readue(bs);
	nalu_debug("num_ref_frames:%ld", res->num_ref_frames);
	res->gaps_in_frame_num_value_allowed_flag = bs_read1bit(bs);
	nalu_debug("gaps_in_frame_num_value_allowed_flag:%d", res->gaps_in_frame_num_value_allowed_flag);
	res->pic_width_in_mbs_minus1 = bs_readue(bs);
	nalu_debug("pic_width_in_mbs_minus1:%ld", res->pic_width_in_mbs_minus1);
	res->pic_height_in_map_units_minus1 = bs_readue(bs);
	nalu_debug("pic_height_in_map_units_minus1:%ld", res->pic_height_in_map_units_minus1);
	res->frame_mbs_only_flag = bs_read1bit(bs);
	nalu_debug("frame_mbs_only_flag:%d", res->frame_mbs_only_flag);
	if (!(res->frame_mbs_only_flag))
	{
		res->mb_adaptive_frame_field_flag = bs_read1bit(bs);
		nalu_debug("mb_adaptive_frame_field_flag:%d", res->mb_adaptive_frame_field_flag);
	}
	res->direct_8x8_inference_flag = bs_read1bit(bs);
	nalu_debug("direct_8x8_inference_flag:%d", res->direct_8x8_inference_flag);
	res->frame_cropping_flag = bs_read1bit(bs);
	nalu_debug("frame_cropping_flag:%d", res->frame_cropping_flag);
	if (res->frame_cropping_flag)
	{
		res->frame_crop_left_offset = bs_readue(bs);
		nalu_debug("frame_crop_left_offset:%ld", res->frame_crop_left_offset);
		res->frame_crop_right_offset = bs_readue(bs);
		nalu_debug("frame_crop_right_offset:%ld", res->frame_crop_right_offset);
		res->frame_crop_top_offset = bs_readue(bs);
		nalu_debug("frame_crop_top_offset:%ld", res->frame_crop_top_offset);
		res->frame_crop_bottom_offset = bs_readue(bs);
		nalu_debug("frame_crop_bottom_offset:%ld", res->frame_crop_bottom_offset);
	}
	res->vui_parameters_present_flag = bs_read1bit(bs);
	nalu_debug("vui_parameters_present_flag:%d", res->vui_parameters_present_flag);
	if (res->vui_parameters_present_flag)
	{
		nalu_vui(&(res->vui), bs);
	}
	rbsp_trailing_bits(bs);

	return res;
}


void nalu_sps_deinit(nalu_sps_t *sps)
{
	if (sps == NULL)
		return;

	nalu_hrd_t *hrd1 = &(sps->vui.hrd1);
	nalu_hrd_t *hrd2 = &(sps->vui.hrd2);
	if (hrd1->bit_rate_value_minus1 != NULL)
		s264_free(hrd1->bit_rate_value_minus1);
	
		if (hrd1->cpb_size_value_minus1 != NULL)
		s264_free(hrd1->cpb_size_value_minus1);
	if (hrd1->cbr_flag != NULL)
		s264_free(hrd1->cbr_flag);

	if (hrd2->bit_rate_value_minus1 != NULL)
		s264_free(hrd2->bit_rate_value_minus1);
	if (hrd2->cpb_size_value_minus1 != NULL)
		s264_free(hrd2->cpb_size_value_minus1);
	if (hrd2->cbr_flag != NULL)
		s264_free(hrd2->cbr_flag);
	
	if (sps->offset_for_ref_frame != NULL)
		s264_free(sps->offset_for_ref_frame);
	s264_free(sps);
}


nalu_pps_t *nalu_parse_pps(bs_t *bs)
{
	
}

void nalu_deinit(void *ptr, u8 type)
{
	if (ptr == NULL)
		return;
	if (type == nalu_type_sps)
	{
		nalu_sps_deinit((nalu_sps_t *) ptr);
	}
}